from typing import Any

import dotenv
from langchain_core.messages import ToolCall, AIMessage, ToolMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
class CustmToolException(Exception):
    '''自定义工具错误异常'''
    def __init__(self, tool_call: ToolCall, exception: Exception) -> None:
        super().__init__()
        self.tool_call= tool_call
        self.exception= exception

@tool
def complex_tool(int_arg: int, float_arg: float, dict_arg: dict):
    '''使用复杂工具进行复杂计算操作'''
    print(dict_arg)
    return int(int_arg * float_arg)

def tool_custom_exception(msg: AIMessage, config: RunnableConfig) -> Any:

    try:
        return complex_tool.invoke(msg.tool_calls[0]['args'], config)
    except Exception as e:
        raise CustmToolException(msg.tool_calls[0], e)

def exception_to_message(inputs: dict) -> dict:
    # 从输入中提取错误信息
    exception= inputs.pop('exception')
    # 将历史消息添加到原始输入中, 以便模型知道它上一次犯了什么错
    messages=[
        AIMessage(content= '', tool_calls= [exception.tool_call]),
        ToolMessage(tool_call_id= exception.tool_call['id'], content= str(exception.exception)),
        HumanMessage(content= '最后一次工具调用引发了异常, 调用参数应该是5, 2.1, {"a": "c"}')
    ]
    inputs['last_output']= messages
    return inputs

# ===== 调用
# 创建prompt, 并预留占位符, 用于存储错误信息
prompt= ChatPromptTemplate.from_messages([
    ('human', '{query}'),
    ('placeholder', '{last_output}')
])

# 创建大模型并绑定工具
llm= ChatOpenAI(model= 'gpt-4o').bind_tools(tools= [complex_tool], tool_choice= 'any')

# 创建链并执行工具
chain= prompt | llm | tool_custom_exception
self_correct_chain= chain.with_fallbacks(
    [exception_to_message | chain], exception_key= 'exception'
)

# 调用自我纠正链完成任务
print(self_correct_chain.invoke({'query': '使用复杂工具, 对应参数为5和2.1'}))